#!/usr/bin/env python3
# -*- coding: utf-8 -*-
import os
import json
import xml.etree.ElementTree as ET
import shutil
from pathlib import Path
from langchain_openai import ChatOpenAI
from dotenv import load_dotenv
import re

load_dotenv()

# 阿里云百炼配置
DASHSCOPE_CONFIG = {
    "api_key": os.getenv('DASHSCOPE_API_KEY'),
    "base_url": os.getenv('BASE_URL'),
    "model": "qwen-long"
}

def clear_folder(folder_path):
    """清空指定文件夹中的所有文件"""
    if folder_path.exists():
        for item in folder_path.iterdir():
            try:
                if item.is_file() or item.is_symlink():
                    item.unlink()
                elif item.is_dir():
                    shutil.rmtree(item)
            except Exception as e:
                print(f'删除 {item} 时出错: {e}')
        print(f'已清空文件夹: {folder_path}')

def split_by_third_level_titles(text: str) -> list:
    """根据三级小标题（如1.3.1格式）分段"""
    try:
        # 初始化大模型
        llm = ChatOpenAI(
            api_key=DASHSCOPE_CONFIG["api_key"],
            base_url=DASHSCOPE_CONFIG["base_url"],
            model=DASHSCOPE_CONFIG["model"],
            temperature=0.1
        )
        
        # 构建提示词
        prompt = f"""请检测以下文本中是否包含三级小标题（形如"1.3.1"、"7.8.3"、"2.1.4"等格式），如果有，请按照这些三级小标题进行分段。

要求：
1. 优先识别三级小标题格式（数字.数字.数字）
2. 按照三级小标题进行分段
3. 绝对不修改原文内容，保持原文完整性
4. 如果没有三级小标题，返回原文本不分段
5. 返回JSON格式，包含segments数组，每个元素有content字段

文本内容：
{text}

请返回JSON格式：
{{
  "segments": [
    {{
      "content": "第一段内容"
    }},
    {{
      "content": "第二段内容"
    }}
  ]
}}"""
        
        # 调用大模型
        response = llm.invoke(prompt)
        response_text = response.content
        
        # 提取JSON内容
        json_match = re.search(r'```json\s*\n?([\s\S]*?)\n?```', response_text, re.IGNORECASE)
        if json_match:
            json_text = json_match.group(1).strip()
        else:
            json_text = response_text.strip()
        
        # 解析JSON
        try:
            result = json.loads(json_text)
            if 'segments' in result and isinstance(result['segments'], list):
                segments = []
                for segment in result['segments']:
                    if 'content' in segment and segment['content'].strip():
                        segments.append(segment['content'].strip())
                # 如果只有一个分段且与原文本相同，说明没有找到三级标题
                if len(segments) == 1 and segments[0] == text.strip():
                    return [text]  # 返回原文本，表示没有分段
                return segments if segments else [text]
            else:
                print("大模型返回格式不正确，使用原文本")
                return [text]
        except json.JSONDecodeError:
            print("大模型返回的JSON格式错误，使用原文本")
            return [text]
            
    except Exception as e:
        print(f"三级标题分段时出错: {e}，使用原文本")
        return [text]

def split_text_with_llm(text: str) -> list:
    """使用大模型对长文本进行语义分段"""
    try:
        # 初始化大模型
        llm = ChatOpenAI(
            api_key=DASHSCOPE_CONFIG["api_key"],
            base_url=DASHSCOPE_CONFIG["base_url"],
            model=DASHSCOPE_CONFIG["model"],
            temperature=0.1
        )
        
        # 构建提示词
        prompt = f"""请对以下文本进行语义分段，要求：
1. 优先检测内部是否有小标题，有小标题的话优先使用小标题分段
2. 没有小标题的话根据语义分段，保证语义完整
3. 每段控制在800字以内
4. 段落之间不重复
5. 绝对不修改原文内容，保持原文完整性
6. 如果分段内部完全是一个整体，可以不分段
7. 返回JSON格式，包含segments数组，每个元素有content字段

文本内容：
{text}

请返回JSON格式：
{{
  "segments": [
    {{
      "content": "第一段内容"
    }},
    {{
      "content": "第二段内容"
    }}
  ]
}}"""
        
        # 调用大模型
        response = llm.invoke(prompt)
        response_text = response.content
        
        # 提取JSON内容
        json_match = re.search(r'```json\s*\n?([\s\S]*?)\n?```', response_text, re.IGNORECASE)
        if json_match:
            json_text = json_match.group(1).strip()
        else:
            json_text = response_text.strip()
        
        # 解析JSON
        try:
            result = json.loads(json_text)
            if 'segments' in result and isinstance(result['segments'], list):
                segments = []
                for segment in result['segments']:
                    if 'content' in segment and segment['content'].strip():
                        segments.append(segment['content'].strip())
                return segments if segments else [text]
            else:
                print("大模型返回格式不正确，使用原文本")
                return [text]
        except json.JSONDecodeError:
            print("大模型返回的JSON格式错误，使用原文本")
            return [text]
            
    except Exception as e:
        print(f"大模型分段时出错: {e}，使用原文本")
        return [text]

def extract_content_from_xml(xml_file_path):
    """从XML文件中提取标题和正文内容，排除<pre>标签中的表格"""
    try:
        tree = ET.parse(xml_file_path)
        root = tree.getroot()
        
        content_sections = []
        current_h1 = None
        current_h2 = None
        current_content = []
        
        # 收集第一个标题之前的内容作为第一段
        pre_title_content = []
        found_first_title = False
        found_catalog = False  # 标记是否找到"目录"
        
        for element in root:
            # 跳过<pre>标签（表格内容）
            if element.tag == 'pre':
                continue
            
            # 处理一级标题 h1（以 *** 开头）
            elif element.tag == 'h1':
                found_first_title = True
                
                # 保存之前的内容段落
                if current_h1 is not None or current_h2 is not None:
                    if current_content:
                        content_sections.append({
                            'h1_title': current_h1,
                            'h2_title': current_h2,
                            'content': '\n'.join(current_content).strip()
                        })
                
                # 开始新的一级标题
                current_h1 = element.text.replace('*** ', '').strip() if element.text else ''
                current_h2 = None
                current_content = []
            
            # 处理二级标题 h2（以 !!! 开头）
            elif element.tag == 'h2':
                found_first_title = True
                
                # 保存之前的内容段落
                if current_h2 is not None and current_content:
                    content_sections.append({
                        'h1_title': current_h1,
                        'h2_title': current_h2,
                        'content': '\n'.join(current_content).strip()
                    })
                
                # 开始新的二级标题
                current_h2 = element.text.replace('!!! ', '').strip() if element.text else ''
                current_content = []
            
            # 处理正文内容 p标签
            elif element.tag == 'p':
                if element.text and element.text.strip():
                    text_content = element.text.strip()
                    # 跳过只包含"!!!"的段落（表格结束标记）
                    if text_content == '!!!':
                        continue
                    
                    if not found_first_title:
                        # 第一个标题之前的内容
                        # 检查是否包含"目录"关键词
                        if '目录' in text_content and not found_catalog:
                            found_catalog = True
                            # 只保存"目录"之前的内容
                            catalog_index = text_content.find('目录')
                            if catalog_index > 0:
                                pre_catalog_content = text_content[:catalog_index].strip()
                                if pre_catalog_content:
                                    pre_title_content.append(pre_catalog_content)
                            # 找到目录后，标记为已找到，但不break，继续处理后续元素
                        elif not found_catalog:
                            # 如果还没找到目录，继续收集内容
                            pre_title_content.append(text_content)
                        # 如果已经找到目录，就不再收集第一段内容，但继续处理后续元素
                    else:
                        # 标题之后的内容
                        current_content.append(text_content)
        
        # 处理最后一个内容段落
        if current_h1 is not None or current_h2 is not None:
            if current_content:
                content_sections.append({
                    'h1_title': current_h1,
                    'h2_title': current_h2,
                    'content': '\n'.join(current_content).strip()
                })
        
        # 如果有第一个标题之前的内容，作为第一段添加
        if pre_title_content:
            content_sections.insert(0, {
                'h1_title': None,
                'h2_title': None,
                'content': '\n'.join(pre_title_content).strip()
            })
        
        return content_sections
    
    except Exception as e:
        print(f"处理文件 {xml_file_path} 时出错: {e}")
        return []

def save_content_to_json(content_sections, json_file_path):
    """将内容段落保存为JSON文件"""
    try:
        # 构建JSON结构
        json_data = {
            'document_sections': []
        }
        
        section_counter = 1
        
        for section in content_sections:
            content_text = section['content']
            # 计算字数（去除空格和换行符）
            word_count = len(content_text.replace(' ', '').replace('\n', '').replace('\r', ''))
            
            # 如果字数超过1000字，进行两级分段处理
            if word_count > 1000:
                print(f"检测到超长文本({word_count}字)，开始分段处理...")
                
                # 第一步：尝试按三级小标题分段
                print("  第一步：检测三级小标题分段...")
                first_level_segments = split_by_third_level_titles(content_text)
                
                final_segments = []
                
                # 第二步：检查每个分段是否仍然超过1000字
                for i, segment in enumerate(first_level_segments):
                    segment_word_count = len(segment.replace(' ', '').replace('\n', '').replace('\r', ''))
                    
                    if segment_word_count > 1000:
                        print(f"    分段{i+1}仍超过1000字({segment_word_count}字)，进行语义分段...")
                        # 使用语义分段进一步分割
                        semantic_segments = split_text_with_llm(segment)
                        final_segments.extend(semantic_segments)
                    else:
                        print(f"    分段{i+1}: {segment_word_count}字，无需进一步分段")
                        final_segments.append(segment)
                
                # 为每个最终分段创建独立的section
                for i, segment in enumerate(final_segments):
                    segment_word_count = len(segment.replace(' ', '').replace('\n', '').replace('\r', ''))
                    section_data = {
                        'section_id': section_counter,
                        'h1_title': section['h1_title'],
                        'h2_title': section['h2_title'],
                        'content': segment,
                        'word_count': segment_word_count,
                        'is_split': True,
                        'split_index': i + 1,
                        'total_splits': len(final_segments),
                        'split_method': 'two_level' if len(first_level_segments) > 1 else 'semantic_only'
                    }
                    json_data['document_sections'].append(section_data)
                    section_counter += 1
                    print(f"  最终分段 {i+1}/{len(final_segments)}: {segment_word_count} 字")
            else:
                # 字数不超过1000字，直接保存
                section_data = {
                    'section_id': section_counter,
                    'h1_title': section['h1_title'],
                    'h2_title': section['h2_title'],
                    'content': content_text,
                    'word_count': word_count,
                    'is_split': False
                }
                json_data['document_sections'].append(section_data)
                section_counter += 1
        
        # 保存为JSON文件
        with open(json_file_path, 'w', encoding='utf-8') as f:
            json.dump(json_data, f, ensure_ascii=False, indent=2)
        
        return True
    
    except Exception as e:
        print(f"保存JSON文件 {json_file_path} 时出错: {e}")
        return False

def batch_convert_xml_to_json():
    """批量转换XML文件到JSON文件"""
    # 定义路径
    xml_folder = Path(r'd:\AAA\LLM_autoQA_API\A_segmenter\xml')
    json_folder = Path(r'd:\AAA\LLM_autoQA_API\A_segmenter\json')
    
    # 清空JSON输出文件夹
    clear_folder(json_folder)
    
    # 创建JSON输出文件夹（如果不存在）
    json_folder.mkdir(exist_ok=True)
    
    # 遍历XML文件夹中的所有XML文件
    xml_files = list(xml_folder.glob('*.xml'))
    
    if not xml_files:
        print("未找到XML文件")
        return
    
    print(f"找到 {len(xml_files)} 个XML文件")
    
    success_count = 0
    
    for xml_file in xml_files:
        print(f"\n处理文件: {xml_file.name}")
        
        # 提取内容
        content_sections = extract_content_from_xml(xml_file)
        
        if content_sections:
            # 生成JSON文件路径（保持原文件名）
            json_file_name = xml_file.stem + '.json'
            json_file_path = json_folder / json_file_name
            
            # 保存到JSON文件
            if save_content_to_json(content_sections, json_file_path):
                success_count += 1
                
                # 统计最终段落数
                with open(json_file_path, 'r', encoding='utf-8') as f:
                    json_data = json.load(f)
                final_sections = len(json_data['document_sections'])
                print(f" {xml_file.name} 最终生成 {final_sections} 个内容段落")
        else:
            print(f"在 {xml_file.name} 中未找到有效内容")
    
    print(f"\n批量转换 {success_count}/{len(xml_files)} 个文件")
    print(f"JSON文件已保存到: {json_folder}")

if __name__ == "__main__":
    # 验证API密钥
    if not DASHSCOPE_CONFIG["api_key"]:
        print("警告：未设置DASHSCOPE_API_KEY环境变量，将跳过大模型分段功能")
    
    batch_convert_xml_to_json()